#ifdef __KERNEL__
#include <linux/errno.h>
#include <linux/string.h>
#else
#include <errno.h>
#include <string.h>
#endif
#include <asm/byteorder.h>
#include "sm2.h"
#include "random.h"

extern struct mpi_ec_ctx sm2_ctx;

/**
 * @brief 计算哈希值Z
 *  Z =SM3(ENTLA ‖ IDA ‖ a ‖ b ‖xG ‖ yG ‖ xA ‖ yA)。
 * @param id     [in] 用户id
 * @param idlen  [in] id长度
 * @param dgst   [out] 哈希值
 * @param public [in] 公钥
 * @return int
 */
static int SM2_compute_id_digest(const unsigned char* id, size_t idlen,
                                 unsigned char dgst[SM3_DIGEST_LENGTH],
                                 EC_POINT* public)
{
    unsigned char idbits[2];
    unsigned char a[SM2_PRIVKEY_LEN];
    unsigned char b[SM2_PRIVKEY_LEN];
    unsigned char xG[SM2_PRIVKEY_LEN];
    unsigned char yG[SM2_PRIVKEY_LEN];
    unsigned char xA[SM2_PRIVKEY_LEN];
    unsigned char yA[SM2_PRIVKEY_LEN];
    unsigned int  buflen = SM2_PRIVKEY_LEN;
    unsigned int  nbytes;
    sm3_ctx_t     ctx[1];
    int           rc = 0;

    mpi_read_buffer(sm2_ctx.a, a, buflen, &nbytes, NULL);
    mpi_read_buffer(sm2_ctx.b, b, buflen, &nbytes, NULL);
    mpi_read_buffer(sm2_ctx.G->x, xG, buflen, &nbytes, NULL);
    mpi_read_buffer(sm2_ctx.G->y, yG, buflen, &nbytes, NULL);
    mpi_read_buffer(public->x, xA, buflen, &nbytes, NULL);
    mpi_read_buffer(public->y, yA, buflen, &nbytes, NULL);

    /* 2-byte id length in bits */
    idbits[0] = ((idlen * 8) >> 8) % 256;
    idbits[1] = (idlen * 8) % 256;

    sm3_init(ctx);
    sm3_update(ctx, idbits, sizeof(idbits));
    sm3_update(ctx, id, idlen);
    sm3_update(ctx, a, sm2_ctx.nbits / 8);
    sm3_update(ctx, b, sm2_ctx.nbits / 8);
    sm3_update(ctx, xG, sm2_ctx.nbits / 8);
    sm3_update(ctx, yG, sm2_ctx.nbits / 8);
    sm3_update(ctx, xA, sm2_ctx.nbits / 8);
    sm3_update(ctx, yA, sm2_ctx.nbits / 8);
    sm3_final(ctx, dgst);
    return rc;
}

/**
 * @brief 预处理
 *      e = SM3(M_) = SM3(Z || M)
 * @param m      [in] 消息
 * @param mlen   [in] 消息长度
 * @param id     [in] 用户id
 * @param idlen  [in] id长度
 * @param dgst   [out] 哈希值
 * @param public [in] 公钥
 * @return int 成功返回0，或错误码
 */
int SM2_compute_message_digest(const unsigned char* m, size_t mlen,
                               const unsigned char* id, size_t idlen,
                               unsigned char dgst[SM3_DIGEST_LENGTH],
                               EC_KEY* public)
{
    int           rc = 0;
    sm3_ctx_t     ctx[1];
    unsigned char ZA[SM3_DIGEST_LENGTH];

    if (!m || !id || !public) return -EINVAL;

    rc = SM2_compute_id_digest(id, idlen, ZA, public->pub);
    if (rc) return rc;

    sm3_init(ctx);
    sm3_update(ctx, ZA, SM3_DIGEST_LENGTH);
    sm3_update(ctx, m, mlen);
    sm3_final(ctx, dgst);

    return rc;
}

/**
 * @brief SM2签名
 *
 * @param dgst      [in] 待签名值
 * @param rs        [out] 签名值
 * @param private   [in] 私钥
 * @return int 成功返回0，或错误码
 */
int SM2_sign(const unsigned char dgst[SM3_DIGEST_LENGTH], SM2_SIG_CTX* rs,
             EC_KEY* private)
{
    struct gcry_mpi_point kG;
    MPI                   k   = NULL;
    MPI                   e   = NULL;
    MPI                   r   = NULL;
    MPI                   s   = NULL;
    MPI                   rk  = NULL;
    MPI                   x1  = NULL;
    MPI                   tmp = NULL;
    MPI                   one = NULL;
    MPI                   d   = NULL;
    unsigned int          nbytes;
    unsigned char         random[SM3_DIGEST_LENGTH];
    int                   rc = -ENOMEM;

    r   = mpi_new(0);
    s   = mpi_new(0);
    rk  = mpi_new(0);
    x1  = mpi_new(0);
    tmp = mpi_new(0);
    if (!r || !s || !rk || !x1 || !tmp) goto free;

    d   = private->priv;
    one = mpi_const(MPI_C_ONE);
    mpi_point_init(&kG);
    mpi_set_ui(one, 1);

    /* A2: e = SM3(M_) */
    e = mpi_read_raw_data(dgst, SM3_DIGEST_LENGTH);
    if (!e) goto free;
loop:
    if (k) mpi_free(k);

    /* A3: 产生随机数k∈[1,n-1] */
    rc = random_generator(random, SM3_DIGEST_LENGTH);
    if (rc) goto free_e;

    k = mpi_read_raw_data(random, SM3_DIGEST_LENGTH);
    if (!k) goto loop;

    /* check: 0 < k < n */
    if (mpi_cmp_ui(k, 1) < 0 || mpi_cmp(k, sm2_ctx.n) > 0) {
        goto loop;
    }

    /* A4: p = (x1,y1) = [k]G */
    mpi_ec_mul_point(&kG, k, sm2_ctx.G, &sm2_ctx);
    if (mpi_ec_get_affine(x1, NULL, &kG, &sm2_ctx)) {
        goto loop;
    }

    /* A5: r = (e + x1) mod n */
    mpi_addm(r, e, x1, sm2_ctx.n);

    /* check: if r == 0 , goto loop */
    if (mpi_cmp_ui(r, 0) == 0) {
        goto loop;
    }

    /* check: if r+k == n, goto loop */
    mpi_add(rk, r, k);
    if (mpi_cmp(rk, sm2_ctx.n) == 0) {
        goto loop;
    }

    /* A6: s = [(1+d)^-1 * (k-r*d) ] mod n */
    mpi_addm(s, d, one, sm2_ctx.n);
    mpi_invm(s, s, sm2_ctx.n);
    mpi_mulm(tmp, r, d, sm2_ctx.n);
    mpi_subm(tmp, k, tmp, sm2_ctx.n);
    mpi_mulm(s, s, tmp, sm2_ctx.n);

    /* s != 0 */
    if (mpi_cmp_ui(s, 0) == 0) {
        mpi_free(k);
        goto loop;
    }

    /* A7: convert into byte string */
    mpi_read_buffer(r, rs->sig_r, SM3_DIGEST_LENGTH, &nbytes, NULL);
    mpi_read_buffer(s, rs->sig_s, SM3_DIGEST_LENGTH, &nbytes, NULL);

    mpi_free(k);
free_e:
    mpi_free(e);
free:
    mpi_free(r);
    mpi_free(s);
    mpi_free(rk);
    mpi_free(x1);
    mpi_free(tmp);
    return rc;
}

/**
 * @brief SM2验签
 *
 * @param dgst      [in] 待验签消息
 * @param rs        [in] 签名值
 * @param public    [in] 公钥
 * @return int 成功返回0，或错误码
 */
int SM2_verify(const unsigned char dgst[SM3_DIGEST_LENGTH], SM2_SIG_CTX* rs,
               EC_KEY* public)
{
    struct gcry_mpi_point sG, tP;
    MPI                   r  = NULL;
    MPI                   s  = NULL;
    MPI                   t  = NULL;
    MPI                   x1 = NULL;
    MPI                   y1 = NULL;
    MPI                   e  = NULL;
    int                   rc = -ENOMEM;

    if (!dgst || !rs || !public) return -EINVAL;

    t  = mpi_new(0);
    x1 = mpi_new(0);
    y1 = mpi_new(0);
    r  = mpi_read_raw_data(rs->sig_r, SM3_DIGEST_LENGTH);
    s  = mpi_read_raw_data(rs->sig_s, SM3_DIGEST_LENGTH);
    if (!t || !x1 || !y1 || !r || !s) goto free;

    mpi_point_init(&sG);
    mpi_point_init(&tP);

    /* B1: check: 0 < r < n */
    if (mpi_cmp_ui(r, 1) < 0 || mpi_cmp(r, sm2_ctx.n) > 0) {
        rc = 1;
        goto free;
    }

    /* B2: check: 0 < s < n */
    if (mpi_cmp_ui(s, 1) < 0 || mpi_cmp(s, sm2_ctx.n) > 0) {
        rc = 2;
        goto free;
    }

    /* B4: e = SM3(M_) */
    e = mpi_read_raw_data(dgst, SM3_DIGEST_LENGTH);
    if (!e) {
        rc = 3;
        goto deinit;
    }

    /* B5: t = (r+s) mod n, t == 0 */
    mpi_addm(t, r, s, sm2_ctx.n);
    if (mpi_cmp_ui(t, 0) == 0) {
        rc = 3;
        goto out;
    }

    /* B6: (x1, y2) = sG + tP */
    mpi_ec_mul_point(&sG, s, sm2_ctx.G, &sm2_ctx);
    mpi_ec_mul_point(&tP, t, public->pub, &sm2_ctx);
    mpi_ec_add_points(&sG, &sG, &tP, &sm2_ctx);
    if (mpi_ec_get_affine(x1, y1, &sG, &sm2_ctx)) {
        rc = 4;
        goto out;
    }

    /* B7: R=(e'+x1')modn */
    mpi_addm(t, e, x1, sm2_ctx.n);

    /* check:R = r */
    if (mpi_cmp(t, r)) {
        rc = 5;
    } else {
        rc = 0;
    }

out:
    mpi_free(e);
deinit:
    mpi_point_free_parts(&sG);
    mpi_point_free_parts(&tP);
free:
    mpi_free(s);
    mpi_free(r);
    mpi_free(y1);
    mpi_free(x1);
    mpi_free(t);
    return rc;
}

/**
 * @brief 密钥派生函数
 *
 * @param z     [in] 输入
 * @param zlen  [in] 输入长度
 * @param k     [out] 密钥
 * @param klen  [in] 要生成的密钥长度
 */
static void KDF(unsigned char* z, const size_t zlen, unsigned char* k,
                const size_t klen)
{
    size_t        i;
    uint32_t      ct = 0x00000001;
    uint32_t      ct_be;
    sm3_ctx_t     ctx[1];
    unsigned char tmp[SM3_DIGEST_LENGTH];
    int           loop = klen / SM3_DIGEST_LENGTH;

    for (i = 0; i < loop; i++) {
        ct_be = __cpu_to_be32(ct);
        sm3_init(ctx);
        sm3_update(ctx, z, zlen);
        sm3_update(ctx, (const unsigned char*)&ct_be, sizeof(ct_be));
        sm3_final(ctx, k);
        k += SM3_DIGEST_LENGTH;
        ++ct;
    }

    loop = klen % SM3_DIGEST_LENGTH;
    if (loop) {
        ct_be = __cpu_to_be32(ct);
        sm3_init(ctx);
        sm3_update(ctx, z, zlen);
        sm3_update(ctx, (const unsigned char*)&ct_be, sizeof(ct_be));
        sm3_final(ctx, tmp);
        memcpy(k, tmp, loop);
    }
}

/**
 * @brief SM2加密函数
 *
 * @param in        [in] 待加密消息
 * @param inlen     [in] 消息长度
 * @param out       [out] 密文
 *  密文结构：[C1 || C2 || C3]
 *      C1: 64字节椭圆曲线点(x1, y1)
 *      C2: SM3(x2||M||y2)
 *      C3: 密文
 * @param outlen    [out] 密文长度
 * @param public    [in] 公钥
 * @return int 成功返回0，或错误码
 */
int SM2_encrypt(const unsigned char* in, size_t inlen, unsigned char* out,
                size_t* outlen, EC_KEY* public)
{
    sm3_ctx_t      sm3_ctx[1];
    unsigned char  x2y2[SM2_PUBKEY_LEN];
    unsigned char  random[SM2_PRIVKEY_LEN];
    unsigned char* C1 = NULL;
    unsigned char* C3 = NULL;
    unsigned char* C2 = NULL;
    MPI_POINT      kG = NULL;
    MPI_POINT      kP = NULL;
    MPI_POINT      S  = NULL;
    MPI            k  = NULL;
    MPI            x  = NULL;
    MPI            y  = NULL;
    MPI            h  = NULL;
    unsigned int   nread;
    size_t         i;
    int            rc = -ENOMEM;

    if (!in || !public) return -EINVAL;
    /* only get the @outlen */
    if (!out && outlen) {
        *outlen = SM2_PUBKEY_LEN + SM3_DIGEST_LENGTH + inlen;
        return 0;
    }

    kG = mpi_point_new(0);
    kP = mpi_point_new(0);
    S  = mpi_point_new(0);
    if (!kG || !kP || !S) goto free_p;

    x = mpi_new(0);
    y = mpi_new(0);
    h = mpi_new(0);
    if (!x || !y || !h) goto free;

    C1 = out;                   // (x1, y1)
    C3 = out + SM2_PUBKEY_LEN;  // Hash (x2||M||y2);
    C2 = out + SM2_PUBKEY_LEN + SM3_DIGEST_LENGTH;
    mpi_set_ui(h, sm2_ctx.h);

loop:
    if (k) mpi_free(k);
    /* A1:产生随机数k∈[1,n-1]; */
    rc = random_generator(random, SM2_PRIVKEY_LEN);
    if (rc) goto free;
    k = mpi_read_raw_data(random, SM2_PRIVKEY_LEN);
    if (mpi_cmp_ui(k, 1) < 0 || mpi_cmp(k, sm2_ctx.n) > 0) {
        goto loop;
    }

    /* A2:C1=[k]G=(x1,y1) */
    mpi_ec_mul_point(kG, k, sm2_ctx.G, &sm2_ctx);
    if (mpi_ec_get_affine(x, y, kG, &sm2_ctx)) {
        goto loop;
    }
    mpi_read_buffer(x, C1, SM2_PRIVKEY_LEN, &nread, NULL);
    mpi_read_buffer(y, C1 + SM2_PRIVKEY_LEN, SM2_PRIVKEY_LEN, &nread, NULL);

    /* A4::计算椭圆曲线点[k]Pb=(x2,y2) */
    mpi_ec_mul_point(kP, k, public->pub, &sm2_ctx);
    if (mpi_ec_get_affine(x, y, kP, &sm2_ctx)) {
        goto loop;
    }
    mpi_read_buffer(x, x2y2, SM2_PRIVKEY_LEN, &nread, NULL);
    mpi_read_buffer(y, x2y2 + SM2_PRIVKEY_LEN, SM2_PRIVKEY_LEN, &nread, NULL);

    /* A3:计算椭圆曲线点S=[h]Pb */
    mpi_ec_mul_point(S, h, public->pub, &sm2_ctx);
    if (mpi_cmp_ui(S->z, 0) == 0) {
        goto free;
    }

    /* A5:t = KDF(x2||y2,klen) */
    KDF(x2y2, SM2_PUBKEY_LEN, C2, inlen);

    /* A6:计算C2=M⊕t; */
    for (i = 0; i < inlen; i++) {
        C2[i] ^= in[i];
    }

    /* A7:计算C3=Hash (x2||M||y2) */
    sm3_init(sm3_ctx);
    sm3_update(sm3_ctx, x2y2, SM2_PRIVKEY_LEN);
    sm3_update(sm3_ctx, in, inlen);
    sm3_update(sm3_ctx, x2y2 + SM2_PRIVKEY_LEN, SM2_PRIVKEY_LEN);
    sm3_final(sm3_ctx, C3);

    rc = 0;
    if (outlen) *outlen = SM2_PUBKEY_LEN + SM3_DIGEST_LENGTH + inlen;
free:
    mpi_free(h);
    mpi_free(k);
    mpi_free(x);
    mpi_free(y);
free_p:
    mpi_point_release(kG);
    mpi_point_release(kP);
    mpi_point_release(S);
    return rc;
}

/**
 * @brief SM2解密函数
 *
 * @param in        [in] 密文
 * @param inlen     [in] 密文长度
 * @param out       [out] 明文
 * @param outlen    [out] 明文长度
 * @param private   [in] 私钥
 * @return int 成功返回0，或错误码
 */
int SM2_decrypt(const unsigned char* in, size_t inlen, unsigned char* out,
                size_t* outlen, EC_KEY* private)
{
    const unsigned char* C1      = NULL;
    const unsigned char* C3      = NULL;
    const unsigned char* C2      = NULL;
    MPI                  x       = NULL;
    MPI                  y       = NULL;
    MPI                  h       = NULL;
    MPI_POINT            S       = NULL;
    MPI_POINT            dC1     = NULL;
    MPI_POINT            xypoint = NULL;
    unsigned char        x2y2[SM2_PUBKEY_LEN];
    sm3_ctx_t            sm3_ctx[1];
    unsigned int         nread;
    size_t               klen  = 0;
    size_t               count = 0;
    size_t               i;
    int                  rc = -ENOMEM;

    if (!in || inlen < SM2_PUBKEY_LEN + SM3_DIGEST_LENGTH || !private) {
        return -EINVAL;
    }
    /* only get the @outlen */
    if (!out && outlen) {
        *outlen = inlen - SM2_PUBKEY_LEN - SM3_DIGEST_LENGTH;
        return 0;
    }

    S       = mpi_point_new(0);
    dC1     = mpi_point_new(0);
    xypoint = mpi_point_new(0);
    if (!S || !dC1 || !xypoint) {
        goto free_p;
    }

    C1   = in;
    C3   = in + SM2_PUBKEY_LEN;
    C2   = in + SM2_PUBKEY_LEN + SM3_DIGEST_LENGTH;
    klen = inlen - SM2_PUBKEY_LEN - SM3_DIGEST_LENGTH;
    h    = mpi_new(0);

    /* B0: 从C1读取(x1, y1) */
    x = mpi_read_raw_data(C1, SM2_PRIVKEY_LEN);
    y = mpi_read_raw_data(C1 + SM2_PRIVKEY_LEN, SM2_PRIVKEY_LEN);
    if (!x || !y || !h) goto free;

    mpi_set(xypoint->x, x);
    mpi_set(xypoint->y, y);
    mpi_set_ui(xypoint->z, 1);
    mpi_set_ui(h, sm2_ctx.h);

    /* B1: 验证C1 = (x1, y1)是否满足椭圆曲线方程 */
    if (!mpi_ec_curve_point(xypoint, &sm2_ctx)) {
        rc = 1;
        goto free;
    }

    /* B2: 计算椭圆曲线点S= [h]C1,若S 是无穷远点,则报错并退出 */
    mpi_ec_mul_point(S, h, xypoint, &sm2_ctx);
    if (mpi_cmp_ui(S->z, 0) == 0) {
        rc = 2;
        goto free;
    }

    /* B3: 计算[d]C1=(x2,y2), 将坐标x2、y2的数据类型转换为比特串 */
    mpi_ec_mul_point(dC1, private->priv, xypoint, &sm2_ctx);
    if (mpi_ec_get_affine(x, y, dC1, &sm2_ctx)) {
        rc = 3;
        goto free;
    }

    mpi_read_buffer(x, x2y2, SM2_PRIVKEY_LEN, &nread, NULL);
    mpi_read_buffer(y, x2y2 + SM2_PRIVKEY_LEN, SM2_PRIVKEY_LEN, &nread, NULL);

    /* B4: 计算t= KDF(x2||y2,klen),若t为全0比特串,则报错并退出 */
    KDF(x2y2, SM2_PUBKEY_LEN, out, klen);
    for (i = 0; i < klen; i++) {
        if (out[i] == 0) {
            ++count;
        }
    }
    if (count == klen) {
        rc = 4;
        goto free;
    }

    /* B5: 从C 中取出比特串C2,计算 M =C2⊕t */
    for (i = 0; i < klen; i++) {
        out[i] ^= C2[i];
    }

    /* B6: 计算u=Hash(x2||M'||y2),从C 中取出比特串C3,若u≠C3,则报错并退出 */
    sm3_init(sm3_ctx);
    sm3_update(sm3_ctx, x2y2, SM2_PRIVKEY_LEN);
    sm3_update(sm3_ctx, out, klen);
    sm3_update(sm3_ctx, x2y2 + SM2_PRIVKEY_LEN, SM2_PRIVKEY_LEN);
    sm3_final(sm3_ctx, x2y2);

    /* check: u == C3 */
    if (memcmp(x2y2, C3, SM3_DIGEST_LENGTH) != 0) {
        rc = 5;
    } else {
        rc      = 0;
        *outlen = klen;
    }
free:
    mpi_free(x);
    mpi_free(y);
    mpi_free(h);
free_p:
    mpi_point_release(xypoint);
    mpi_point_release(dC1);
    mpi_point_release(S);
    return rc;
}

/**
 * @brief SM2密钥交换初始化函数
 *      1.对id进行预处理;2.计算2^w
 * @param ctx           [in/out] 上下文
 * @param key           [in] 本端密钥对
 * @param peer_pubkey   [in] 对端公钥
 * @param id            [in] 本端用户id
 * @param idlen         [in] id长度
 * @param rid           [in] 对端用户id
 * @param ridlen        [in] id长度
 * @param is_initiator  [in] 是不是发起者
 * @param do_checksum   [in] 是否进行校验
 * @return int 成功返回0，或错误码
 */
int SM2_KAP_CTX_init(SM2_KAP_CTX* ctx, EC_KEY* key, EC_POINT* peer_pubkey,
                     const char* id, size_t idlen, const char* rid,
                     size_t ridlen, int is_initiator, int do_checksum)
{
    int rc = -ENOMEM;

    if (!ctx || !key || !peer_pubkey || !id || !rid) return -EINVAL;

    ctx->key       = ec_key_new();
    ctx->two_pow_w = mpi_new(0);
    ctx->t         = mpi_new(0);
    if (!ctx->key || !ctx->two_pow_w || !ctx->t) goto free;

    mpi_point_init(&ctx->rpub);
    mpi_point_init(&ctx->R);
    ctx->do_checksum  = do_checksum;
    ctx->is_initiator = is_initiator;

    mpi_set(ctx->key->priv, key->priv);
    mpi_set(ctx->key->pub->x, key->pub->x);
    mpi_set(ctx->key->pub->y, key->pub->y);
    mpi_set(ctx->key->pub->z, key->pub->z);
    mpi_set(ctx->rpub.x, peer_pubkey->x);
    mpi_set(ctx->rpub.y, peer_pubkey->y);
    mpi_set(ctx->rpub.z, peer_pubkey->z);

    /* 计算Za, Zb */
    rc = SM2_compute_id_digest((const unsigned char*)id, idlen, ctx->Z,
                               ctx->key->pub);
    if (rc) goto free;
    rc = SM2_compute_id_digest((const unsigned char*)rid, ridlen, ctx->rZ,
                               &ctx->rpub);
    if (rc) goto free;
    /* 计算 w = ⌈(⌈log2(n)⌉/2)⌉-1 = ceil(nbits / 2) - 1 = 256/2 - 1 = 127 */
    mpi_set_ui(ctx->two_pow_w, 1);
    mpi_lshift(ctx->two_pow_w, ctx->two_pow_w, 256 / 2 - 1);
    return rc;
free:
    mpi_free(ctx->t);
    mpi_free(ctx->two_pow_w);
    ec_key_free(ctx->key);
    mpi_point_free_parts(&ctx->rpub);
    mpi_point_free_parts(&ctx->R);
    return rc;
}

void SM2_KAP_CTX_cleanup(SM2_KAP_CTX* ctx)
{
    if (!ctx) return;
    mpi_free(ctx->t);
    mpi_free(ctx->two_pow_w);
    ec_key_free(ctx->key);
    mpi_point_free_parts(&ctx->rpub);
    mpi_point_free_parts(&ctx->R);
}

/**
 * @brief SM2密钥交换准备函数
 *
 * @param ctx [in/out] 上下文
 * @param R   [out] 曲线点
 * @return int 成功返回0，或错误码
 */
int SM2_KAP_prepare(SM2_KAP_CTX* ctx, unsigned char R[SM2_PUBKEY_LEN])
{
    int           rc = -ENOMEM;
    MPI           r  = NULL;
    MPI           x_ = NULL;
    MPI_POINT     rG = NULL;
    unsigned int  nread;
    unsigned char random[SM2_PRIVKEY_LEN];

    if (!ctx || !R) return -EINVAL;

    x_ = mpi_new(0);
    rG = mpi_point_new(0);
    if (!rG || !x_) goto free;

loop:
    /* 第一步 产生随机数 rA∈[1,n-1]; */
    rc = random_generator(random, SM2_PRIVKEY_LEN);
    if (rc) goto free;
    r = mpi_read_raw_data(random, SM2_PRIVKEY_LEN);
    if (!r) goto free;
    if (mpi_cmp_ui(r, 1) < 0 || mpi_cmp(r, sm2_ctx.n) > 0) {
        mpi_free(r);
        goto loop;
    }

    /* 第二步 生成rG 计算椭圆曲线点RA= [rA]G=(x1,y1); */
    mpi_ec_mul_point(rG, r, sm2_ctx.G, &sm2_ctx);
    if (mpi_ec_get_affine(ctx->R.x, ctx->R.y, rG, &sm2_ctx)) {
        rc = 1;
        goto free;
    }
    mpi_set(ctx->R.z, rG->z);

    /* 第四步 从RA中取出域元素x1 计算
        x_ = 2^w + (x1 & (2^w -1)) = 2^w + (x mod 2^w) */
    mpi_mod(x_, ctx->R.x, ctx->two_pow_w);
    mpi_add(x_, ctx->two_pow_w, x_);

    /* 第五步 计算 t=(d + x_·rA)mod n; */
    mpi_mulm(ctx->t, x_, r, sm2_ctx.n);
    mpi_addm(ctx->t, ctx->t, ctx->key->priv, sm2_ctx.n);

    /* 第三步 输出R 将RA发送给用户B; */
    mpi_read_buffer(ctx->R.x, R, SM2_PRIVKEY_LEN, &nread, NULL);
    mpi_read_buffer(ctx->R.y, R + SM2_PRIVKEY_LEN, SM2_PRIVKEY_LEN, &nread,
                    NULL);
free:
    mpi_free(r);
    mpi_free(x_);
    mpi_point_release(rG);
    return rc;
}

/**
 * @brief 计算校验值
 *  S=SM3(flage‖yU‖SM3(xU‖ZA‖ZB‖x1‖y1‖x2‖y2)
 * @param flage [in] 校验标志 0x02或0x03
 * @param Ux    [in] 椭圆曲线点x
 * @param Uy    [in] 椭圆曲线点y
 * @param Z     [in] 发起端Z
 * @param rZ    [in] 响应端Z
 * @param p1    [in] 发起端椭圆曲线点
 * @param p2    [in] 响应端椭圆曲线点
 * @param S     [out] 校验值
 */
static void SM3_S(char flage, MPI Ux, MPI Uy, unsigned char* Z,
                  unsigned char* rZ, MPI_POINT p1, MPI_POINT p2,
                  unsigned char* S)
{
    sm3_ctx_t     ctx[1];
    unsigned int  nread;
    unsigned char yu[SM2_PRIVKEY_LEN];
    unsigned char tmp[SM3_DIGEST_LENGTH];

    sm3_init(ctx);

    mpi_read_buffer(Ux, tmp, SM2_PRIVKEY_LEN, &nread, NULL);
    sm3_update(ctx, tmp, SM2_PRIVKEY_LEN);

    sm3_update(ctx, Z, SM3_DIGEST_LENGTH);
    sm3_update(ctx, rZ, SM3_DIGEST_LENGTH);

    mpi_read_buffer(p1->x, tmp, SM2_PRIVKEY_LEN, &nread, NULL);
    sm3_update(ctx, tmp, SM2_PRIVKEY_LEN);

    mpi_read_buffer(p1->y, tmp, SM2_PRIVKEY_LEN, &nread, NULL);
    sm3_update(ctx, tmp, SM2_PRIVKEY_LEN);

    mpi_read_buffer(p2->x, tmp, SM2_PRIVKEY_LEN, &nread, NULL);
    sm3_update(ctx, tmp, SM2_PRIVKEY_LEN);

    mpi_read_buffer(p2->y, tmp, SM2_PRIVKEY_LEN, &nread, NULL);
    sm3_update(ctx, tmp, SM2_PRIVKEY_LEN);

    sm3_final(ctx, tmp);

    mpi_read_buffer(Uy, yu, SM2_PRIVKEY_LEN, &nread, NULL);

    sm3_init(ctx);
    sm3_update(ctx, (const unsigned char*)&flage, 1);
    sm3_update(ctx, yu, SM2_PRIVKEY_LEN);
    sm3_update(ctx, tmp, SM3_DIGEST_LENGTH);
    sm3_final(ctx, S);
}

/**
 * @brief 计算交换密钥
 *
 * @param ctx       [in/out] 上下文
 * @param R         [in] 对端曲线点
 * @param key       [out] 密钥
 * @param keylen    [in] 要生成的密钥长度
 * @param checksum  [out] 校验值（可选）
 * @return int 成功返回0，或错误码
 */
int SM2_KAP_compute_key(SM2_KAP_CTX* ctx, unsigned char R[SM2_PUBKEY_LEN],
                        unsigned char* key, size_t keylen,
                        unsigned char checksum[SM3_DIGEST_LENGTH])
{
    int           rc  = -ENOMEM;
    MPI           y   = NULL;
    MPI           x   = NULL;
    MPI           xp_ = NULL;
    MPI           h   = NULL;
    MPI_POINT     U   = NULL;
    MPI_POINT     Rp  = NULL;
    unsigned int  nread;
    unsigned char kdfbuf[SM2_PUBKEY_LEN + SM3_DIGEST_LENGTH * 2];

    if (!ctx || !R || !key) return -EINVAL;

    xp_ = mpi_new(0);
    U   = mpi_point_new(0);
    Rp  = mpi_point_new(0);
    h   = mpi_alloc_set_ui(sm2_ctx.h);
    x   = mpi_read_raw_data(R, SM2_PRIVKEY_LEN);
    y   = mpi_read_raw_data(R + SM2_PRIVKEY_LEN, SM2_PRIVKEY_LEN);
    if (!xp_ || !U || !Rp || !h || !x || !y) goto free;

    mpi_set(Rp->x, x);
    mpi_set(Rp->y, y);
    mpi_set_ui(Rp->z, 1);

    /* 第六步：验证RB是否满足椭圆曲线方程 */
    if (!mpi_ec_curve_point(Rp, &sm2_ctx)) {
        rc = 1;
        goto free;
    }

    /* 第六步：从RB中取出域元素x2, 计算
        x_ = 2^w + (x2 and (2^w - 1)) = 2^w + (x2 mod 2^w) */
    mpi_mod(xp_, Rp->x, ctx->two_pow_w);
    mpi_add(xp_, ctx->two_pow_w, xp_);

    /* 第七步：U = ht * (P + x_ * R)=(xU ,yU ), check U != O */
    mpi_ec_mul_point(U, xp_, Rp, &sm2_ctx);
    mpi_ec_add_points(U, &ctx->rpub, U, &sm2_ctx);
    mpi_ec_mul_point(U, ctx->t, U, &sm2_ctx);
    mpi_ec_mul_point(U, h, U, &sm2_ctx);
    if (mpi_cmp_ui(U->z, 0) == 0) {
        rc = 2;
        goto free;
    }
    if (mpi_ec_get_affine(x, y, U, &sm2_ctx)) {
        rc = 3;
        goto free;
    }

    /* 第八步: key = KDF(xU ‖yU ‖ZA‖ZB,klen) */
    mpi_read_buffer(x, kdfbuf, SM2_PRIVKEY_LEN, &nread, NULL);
    mpi_read_buffer(y, kdfbuf + SM2_PRIVKEY_LEN, SM2_PRIVKEY_LEN, &nread, NULL);
    if (ctx->is_initiator) {
        memcpy(kdfbuf + SM2_PUBKEY_LEN, ctx->Z, SM3_DIGEST_LENGTH);
        memcpy(kdfbuf + SM2_PUBKEY_LEN + SM3_DIGEST_LENGTH, ctx->rZ,
               SM3_DIGEST_LENGTH);
    } else {
        memcpy(kdfbuf + SM2_PUBKEY_LEN, ctx->rZ, SM3_DIGEST_LENGTH);
        memcpy(kdfbuf + SM2_PUBKEY_LEN + SM3_DIGEST_LENGTH, ctx->Z,
               SM3_DIGEST_LENGTH);
    }
    KDF(kdfbuf, SM2_PUBKEY_LEN + SM3_DIGEST_LENGTH * 2, key, keylen);

    /* 第九步（可选）：SM3(0x02‖yU‖SM3(xU‖ZA‖ZB‖x1‖y1‖x2‖y2)
       第十步（可选）：SM3(0x03‖yU‖SM3(xU‖ZA‖ZB‖x1‖y1‖x2‖y2) */
    if (ctx->do_checksum && checksum) {
        if (ctx->is_initiator) {
            SM3_S(0x02, x, y, ctx->Z, ctx->rZ, &ctx->R, Rp, ctx->checksum);
            SM3_S(0x03, x, y, ctx->Z, ctx->rZ, &ctx->R, Rp, checksum);
        } else {
            SM3_S(0x03, x, y, ctx->rZ, ctx->Z, Rp, &ctx->R, ctx->checksum);
            SM3_S(0x02, x, y, ctx->rZ, ctx->Z, Rp, &ctx->R, checksum);
        }
    }
    rc = 0;
free:
    mpi_free(h);
    mpi_free(x);
    mpi_free(y);
    mpi_free(xp_);
    mpi_point_release(U);
    mpi_point_release(Rp);
    return rc;
}

/**
 * @brief 校验
 *
 * @param ctx       [in/out] 上下文
 * @param checksum  [in] 对端校验值
 * @return int 成功返回0，或非0
 */
int SM2_KAP_final_check(SM2_KAP_CTX*        ctx,
                        const unsigned char checksum[SM3_DIGEST_LENGTH])
{
    if (!ctx || !checksum) return -EINVAL;
    if (ctx->do_checksum) {
        if (memcmp(ctx->checksum, checksum, SM3_DIGEST_LENGTH)) {
            return 1;
        }
    }
    return 0;
}